"""
# @Time    : 2021/7/1 6:52 下午
# @Author  : hezhiqiang01
# @Email   : hezhiqiang01@baidu.com
# @File    : r_mappo.py
"""

import numpy as np
import torch
import torch.nn as nn
from utils.util import get_gard_norm, huber_loss, mse_loss
from utils.valuenorm import ValueNorm
from algorithms.utils.util import check


class RMAPPO():
    """
    MAPPO训练器类，用于更新策略。
    :param args: (argparse.Namespace) 包含相关模型、策略和环境信息的参数。
    :param policy: (R_MAPPO_Policy) 要更新的策略。
    :param device: (torch.device) 指定运行设备(cpu/gpu)。
    """

    def __init__(self,
                 args,
                 policy,
                 device=torch.device("cpu")):

        self.device = device
        self.tpdv = dict(dtype=torch.float32, device=device)
        self.policy = policy

        self.clip_param = args.clip_param  # PPO裁剪参数
        self.ppo_epoch = args.ppo_epoch  # PPO训练轮数
        self.num_mini_batch = args.num_mini_batch  # 小批量数量
        self.data_chunk_length = args.data_chunk_length  # 数据块长度
        self.value_loss_coef = args.value_loss_coef  # 价值损失系数
        self.entropy_coef = args.entropy_coef  # 熵系数
        self.max_grad_norm = args.max_grad_norm  # 最大梯度范数
        self.huber_delta = args.huber_delta  # Huber损失的delta参数

        # 各种功能开关
        self._use_recurrent_policy = args.use_recurrent_policy  # 是否使用循环策略
        self._use_naive_recurrent = args.use_naive_recurrent_policy  # 是否使用简单循环策略
        self._use_max_grad_norm = args.use_max_grad_norm  # 是否使用最大梯度范数
        self._use_clipped_value_loss = args.use_clipped_value_loss  # 是否使用裁剪价值损失
        self._use_huber_loss = args.use_huber_loss  # 是否使用Huber损失
        self._use_popart = args.use_popart  # 是否使用PopArt归一化
        self._use_valuenorm = args.use_valuenorm  # 是否使用价值归一化
        self._use_value_active_masks = args.use_value_active_masks  # 是否在价值损失中使用活动掩码
        self._use_policy_active_masks = args.use_policy_active_masks  # 是否在策略损失中使用活动掩码

        assert (self._use_popart and self._use_valuenorm) == False, (
            "self._use_popart和self._use_valuenorm不能同时设置为True")

        # 根据设置选择价值归一化器
        if self._use_popart:
            self.value_normalizer = self.policy.critic.v_out
        elif self._use_valuenorm:
            self.value_normalizer = ValueNorm(1, device=self.device)
        else:
            self.value_normalizer = None

    def cal_value_loss(self, values, value_preds_batch, return_batch, active_masks_batch):
        """
        计算价值函数损失。
        :param values: (torch.Tensor) 价值函数预测。
        :param value_preds_batch: (torch.Tensor) 来自数据批次的"旧"价值预测（用于价值裁剪损失）。
        :param return_batch: (torch.Tensor) 累积回报。
        :param active_masks_batch: (torch.Tensor) 表示智能体在给定时间步是活动的还是死亡的。

        :return value_loss: (torch.Tensor) 价值函数损失。
        """
        # 计算裁剪后的价值预测
        value_pred_clipped = value_preds_batch + (values - value_preds_batch).clamp(-self.clip_param,
                                                                                    self.clip_param)
        # 根据是否使用归一化计算误差
        if self._use_popart or self._use_valuenorm:
            self.value_normalizer.update(return_batch)
            error_clipped = self.value_normalizer.normalize(return_batch) - value_pred_clipped
            error_original = self.value_normalizer.normalize(return_batch) - values
        else:
            error_clipped = return_batch - value_pred_clipped
            error_original = return_batch - values

        # 根据是否使用Huber损失计算价值损失
        if self._use_huber_loss:
            value_loss_clipped = huber_loss(error_clipped, self.huber_delta)
            value_loss_original = huber_loss(error_original, self.huber_delta)
        else:
            value_loss_clipped = mse_loss(error_clipped)
            value_loss_original = mse_loss(error_original)

        # 根据是否使用裁剪价值损失选择最终损失
        if self._use_clipped_value_loss:
            value_loss = torch.max(value_loss_original, value_loss_clipped)
        else:
            value_loss = value_loss_original

        # 根据是否使用活动掩码计算最终价值损失
        if self._use_value_active_masks:
            value_loss = (value_loss * active_masks_batch).sum() / active_masks_batch.sum()
        else:
            value_loss = value_loss.mean()

        return value_loss

    def ppo_update(self, sample, update_actor=True):
        """
        更新演员和评论家网络。
        :param sample: (Tuple) 包含用于更新网络的数据批次。
        :update_actor: (bool) 是否更新演员网络。

        :return value_loss: (torch.Tensor) 价值函数损失。
        :return critic_grad_norm: (torch.Tensor) 评论家更新的梯度范数。
        :return policy_loss: (torch.Tensor) 演员(策略)损失值。
        :return dist_entropy: (torch.Tensor) 动作熵。
        :return actor_grad_norm: (torch.Tensor) 演员更新的梯度范数。
        :return imp_weights: (torch.Tensor) 重要性采样权重。
        """
        # 解包样本数据
        share_obs_batch, obs_batch, rnn_states_batch, rnn_states_critic_batch, actions_batch, \
        value_preds_batch, return_batch, masks_batch, active_masks_batch, old_action_log_probs_batch, \
        adv_targ, available_actions_batch = sample

        # 将数据转换为张量并移动到指定设备
        old_action_log_probs_batch = check(old_action_log_probs_batch).to(**self.tpdv)
        adv_targ = check(adv_targ).to(**self.tpdv)
        value_preds_batch = check(value_preds_batch).to(**self.tpdv)
        return_batch = check(return_batch).to(**self.tpdv)
        active_masks_batch = check(active_masks_batch).to(**self.tpdv)

        # 重塑以在单个前向传递中完成所有步骤
        values, action_log_probs, dist_entropy = self.policy.evaluate_actions(share_obs_batch,
                                                                              obs_batch,
                                                                              rnn_states_batch,
                                                                              rnn_states_critic_batch,
                                                                              actions_batch,
                                                                              masks_batch,
                                                                              available_actions_batch,
                                                                              active_masks_batch)
        # 演员更新
        # 计算重要性权重
        imp_weights = torch.exp(action_log_probs - old_action_log_probs_batch)

        # 计算替代目标
        surr1 = imp_weights * adv_targ
        surr2 = torch.clamp(imp_weights, 1.0 - self.clip_param, 1.0 + self.clip_param) * adv_targ

        # 根据是否使用策略活动掩码计算策略动作损失
        if self._use_policy_active_masks:
            policy_action_loss = (-torch.sum(torch.min(surr1, surr2),
                                             dim=-1,
                                             keepdim=True) * active_masks_batch).sum() / active_masks_batch.sum()
        else:
            policy_action_loss = -torch.sum(torch.min(surr1, surr2), dim=-1, keepdim=True).mean()

        policy_loss = policy_action_loss

        # 清零演员优化器的梯度
        self.policy.actor_optimizer.zero_grad()

        # 如果更新演员，计算损失并反向传播
        if update_actor:
            (policy_loss - dist_entropy * self.entropy_coef).backward()

        # 根据是否使用最大梯度范数裁剪梯度
        if self._use_max_grad_norm:
            actor_grad_norm = nn.utils.clip_grad_norm_(self.policy.actor.parameters(), self.max_grad_norm)
        else:
            actor_grad_norm = get_gard_norm(self.policy.actor.parameters())

        # 执行演员优化器的步骤
        self.policy.actor_optimizer.step()

        # 评论家更新
        # 计算价值损失
        value_loss = self.cal_value_loss(values, value_preds_batch, return_batch, active_masks_batch)

        # 清零评论家优化器的梯度
        self.policy.critic_optimizer.zero_grad()

        # 计算损失并反向传播
        (value_loss * self.value_loss_coef).backward()

        # 根据是否使用最大梯度范数裁剪梯度
        if self._use_max_grad_norm:
            critic_grad_norm = nn.utils.clip_grad_norm_(self.policy.critic.parameters(), self.max_grad_norm)
        else:
            critic_grad_norm = get_gard_norm(self.policy.critic.parameters())

        # 执行评论家优化器的步骤
        self.policy.critic_optimizer.step()

        return value_loss, critic_grad_norm, policy_loss, dist_entropy, actor_grad_norm, imp_weights

    def train(self, buffer, update_actor=True):
        """
        使用小批量梯度下降执行训练更新。
        :param buffer: (SharedReplayBuffer) 包含训练数据的缓冲区。
        :param update_actor: (bool) 是否更新演员网络。

        :return train_info: (dict) 包含有关训练更新的信息（例如损失、梯度范数等）。
        """
        # 计算优势
        if self._use_popart or self._use_valuenorm:
            advantages = buffer.returns[:-1] - self.value_normalizer.denormalize(buffer.value_preds[:-1])
        else:
            advantages = buffer.returns[:-1] - buffer.value_preds[:-1]
        
        # 标准化优势
        advantages_copy = advantages.copy()
        advantages_copy[buffer.active_masks[:-1] == 0.0] = np.nan
        mean_advantages = np.nanmean(advantages_copy)
        std_advantages = np.nanstd(advantages_copy)
        advantages = (advantages - mean_advantages) / (std_advantages + 1e-5)

        # 初始化训练信息字典
        train_info = {}

        train_info['value_loss'] = 0
        train_info['policy_loss'] = 0
        train_info['dist_entropy'] = 0
        train_info['actor_grad_norm'] = 0
        train_info['critic_grad_norm'] = 0
        train_info['ratio'] = 0

        # 执行多个PPO训练轮次
        for _ in range(self.ppo_epoch):
            # 根据策略类型选择数据生成器
            if self._use_recurrent_policy:
                data_generator = buffer.recurrent_generator(advantages, self.num_mini_batch, self.data_chunk_length)
            elif self._use_naive_recurrent:
                data_generator = buffer.naive_recurrent_generator(advantages, self.num_mini_batch)
            else:
                data_generator = buffer.feed_forward_generator(advantages, self.num_mini_batch)

            # 对每个样本进行更新
            for sample in data_generator:
                value_loss, critic_grad_norm, policy_loss, dist_entropy, actor_grad_norm, imp_weights \
                    = self.ppo_update(sample, update_actor)

                # 累积训练信息
                train_info['value_loss'] += value_loss.item()
                train_info['policy_loss'] += policy_loss.item()
                train_info['dist_entropy'] += dist_entropy.item()
                train_info['actor_grad_norm'] += actor_grad_norm
                train_info['critic_grad_norm'] += critic_grad_norm
                train_info['ratio'] += imp_weights.mean()

        # 计算更新次数
        num_updates = self.ppo_epoch * self.num_mini_batch

        # 计算平均训练信息
        for k in train_info.keys():
            train_info[k] /= num_updates

        return train_info

    def prep_training(self):
        """
        准备训练模式。
        将策略的演员和评论家网络设置为训练模式。
        """
        self.policy.actor.train()
        self.policy.critic.train()

    def prep_rollout(self):
        """
        准备推理模式。
        将策略的演员和评论家网络设置为评估模式。
        """
        self.policy.actor.eval()
        self.policy.critic.eval()

